package cn.liwq.demo.java.nio.chat;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.charset.StandardCharsets;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ThreadLocalRandom;

/**
 * @author lweiqiang
 * @date 2021/4/26
 */
public class SimpleGroupChat {
    public static void main(String[] args) {
        try {
            Selector selector = Selector.open();
            ServerSocketChannel serverSocketChannel = ServerSocketChannel.open();
            serverSocketChannel.configureBlocking(false);
            serverSocketChannel.bind(new InetSocketAddress(8080));
            serverSocketChannel.register(selector, SelectionKey.OP_ACCEPT);
            while (true) {
                selector.select();
                Iterator<SelectionKey> iterator = selector.selectedKeys().iterator();
                while (iterator.hasNext()) {
                    SelectionKey key = iterator.next();
                    if (key.isAcceptable()) {
                        ServerSocketChannel ssChan = (ServerSocketChannel) key.channel();
                        SocketChannel socketChannel = ssChan.accept();
                        System.out.println("建立新连接：" + socketChannel.getRemoteAddress());
                        socketChannel.configureBlocking(false);
                        socketChannel.register(selector, SelectionKey.OP_READ);
                        ChatHolder.join(socketChannel);
                    } else if (key.isReadable()) {
                        SocketChannel socketChannel = (SocketChannel) key.channel();
                        ByteBuffer byteBuffer = ByteBuffer.allocate(1024);
                        int len = socketChannel.read(byteBuffer);
                        if (len > 0) {
                            String msg = new String(byteBuffer.array()).trim();
                            if (msg.equals("quit")) {
                                ChatHolder.quit(socketChannel);
                                key.cancel();
                                socketChannel.close();
                            } else {
                                ChatHolder.propogate(socketChannel, msg);
                            }
                        }
                    }
                    iterator.remove();
                }
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public static class ChatHolder {
        private static final Map<SocketChannel, String> USER_MAP = new ConcurrentHashMap<>();

        public static void join(SocketChannel socketChannel) {
            String userId = "用户" + ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE);
            for (SocketChannel channel : USER_MAP.keySet()) {
                send(channel, userId + "加入群聊\n\r");
            }
            USER_MAP.put(socketChannel, userId);
        }

        public static void quit(SocketChannel socketChannel) {
            String userId = USER_MAP.remove(socketChannel);
            for (SocketChannel channel : USER_MAP.keySet()) {
                send(channel, userId + "退出了群聊\n\r");
            }
        }

        public static void propogate(SocketChannel socketChannel, String msg) {
            String userId = USER_MAP.get(socketChannel);
            for (SocketChannel channel : USER_MAP.keySet()) {
                if (!channel.equals(socketChannel)) {
                    send(channel, userId + ":" + msg + "\n\r");
                }
            }
        }

        private static void send(SocketChannel socketChannel, String msg) {
            try {
                ByteBuffer writeBuffer = ByteBuffer.allocate(1024);
                writeBuffer.put(msg.getBytes(StandardCharsets.UTF_8));
                writeBuffer.flip();
                socketChannel.write(writeBuffer);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }
}
